﻿#region Copyright information
// 
// Copyright © 2005-2013 Yongkee Cho. All rights reserved.
// 
// This code is a part of the Biological Object Library and governed under the terms of the
// GNU Lesser General  Public License (LGPL) version 2.1 which accompanies this distribution.
// For more information on the LGPL, please visit http://bol.codeplex.com/license.
// 
// - Filename: NaiveBayesClassifier.cs
// - Author: Yongkee Cho
// - Email: yongkeecho@gmail.com
// - Date Created: 2012-12-08 8:26 AM
// - Last Modified: 2013-01-25 3:59 PM
// 
#endregion
using System;
using System.Collections.Generic;
using System.Linq;
using BOL.Linq;
using BOL.Linq.Descriptive;
using BOL.Maths.Distributions;

namespace BOL.Algorithms.Classification
{
    /// <summary>
    /// Implements a generic Naive Bayes classifier.
    /// </summary>
    /// <see>http://en.wikipedia.org/wiki/Naive_Bayes_classifier</see>
    /// <see>http://www.mathworks.com/help/toolbox/stats/naivebayesclass.html</see>
    public class NaiveBayesClassifier<TFeature, TClass> : ClassifierBase<TFeature, TClass>, IEquatable<NaiveBayesClassifier<TFeature, TClass>>
    {
        #region Public properties

        public IDictionary<TClass, double> Prior { get; private set; }

        public IDictionary<TClass, IList<IDistribution<TFeature>>> Likelihood { get; private set; }

        #endregion

        #region Constructor

        public NaiveBayesClassifier(IDictionary<TClass, double> prior, IDictionary<TClass, IList<IDistribution<TFeature>>> likelihood) 
        {
            if (prior == null)
                throw new ArgumentNullException("prior");
            if (likelihood == null)
                throw new ArgumentNullException("likelihood");

            Likelihood = likelihood;
            Prior = prior;
        }

        #endregion
    
        #region Public methods

        public static NaiveBayesClassifier<TKey, TClass> Train<TSource, TKey>(IEnumerable<TSource> source, Func<TSource, TClass> classSelector, Func<TSource, TKey[]> featuresSelector)
            where TKey: struct, IComparable
        {
            if (source == null)
                throw new ArgumentNullException("source");
            if (classSelector == null)
                throw new ArgumentNullException("classSelector");
            if (featuresSelector == null)
                throw new ArgumentNullException("featuresSelector");

            var list = source.GroupBy(classSelector).Select(x => new { Class = x.Key, Features = x.Select(featuresSelector).ToArray(), Count = x.Count() }).ToList();
            var total = list.Sum(x => x.Count);
            var featureLength = list.First().Features.Length;

            // Estimate prior
            var prior = list.ToDictionary(x => x.Class, x => x.Count / (double)total);

            // Estimate likelihood
            var likelihood = list.ToDictionary(x => x.Class, x => MaximumLikelihoodEstimate(new TKey[featureLength].Select(y => new CategoricalDistribution<TKey>(new Dictionary<TKey, double>())), x.Features));

            return new NaiveBayesClassifier<TKey, TClass>(prior, likelihood);
        }

        public static NaiveBayesClassifier<double, TClass> Train<TSource>(IEnumerable<TSource> source, Func<TSource, TClass> classSelector, Func<TSource, double[]> featuresSelector)
        {
            if (source == null)
                throw new ArgumentNullException("source");
            if (classSelector == null)
                throw new ArgumentNullException("classSelector");
            if (featuresSelector == null)
                throw new ArgumentNullException("featuresSelector");

            var list = source.GroupBy(classSelector).Select(x => new { Class = x.Key, Features = x.Select(featuresSelector).ToArray(), Count = x.Count() }).ToList();
            var total = list.Sum(x => x.Count);
            var featureLength = list.First().Features.Length;

            // Estimate prior
            var prior = list.ToDictionary(x => x.Class, x => x.Count / (double)total);

            // Estimate likelihood
            var likelihood = list.ToDictionary(x => x.Class, x => MaximumLikelihoodEstimate(new double[featureLength].Select(y => GaussianDistribution.StandardDistribution), x.Features));

            return new NaiveBayesClassifier<double, TClass>(prior, likelihood);
        }

        private static IList<IDistribution<T>> MaximumLikelihoodEstimate<T>(IEnumerable<IDistribution<T>> likelihood, IEnumerable<T[]> data)
        {
            if (likelihood == null)
                throw new ArgumentNullException("likelihood");
            if (data == null)
                throw new ArgumentNullException("data");
            
            // swap row and column
            // convert IEnumerable<IList<TFeature>> into IList<IEnumerable<TFeature>>
            var list = data.ToList();
            var numberOfFeatures = list.First().Length;
            var inverted = likelihood.Zip(Enumerable.Range(0, numberOfFeatures).Select(i => list.Select(x => x[i])),
                                          (l, i) => new {Distribution = l, Data = i}).ToList();
            inverted.ForEach(x => x.Distribution.MaximumLikelihoodEstimate(x.Data));

            return inverted.Select(x => x.Distribution).ToList();
        }

        public override TClass Classify(TFeature[] features)
        {
            if (features.Length != Likelihood.Count)
                throw new ArgumentException("Dimension mismatches.");

            var logPosterior = Prior.ToDictionary(x => x.Key, x => Math.Log(x.Value) + 
                Likelihood[x.Key].Zip(features, (l, f) => Math.Log(l.Pdf(f))).Sum());

            return logPosterior.ArgMax();
        }

        #endregion

        #region IEquatable<NaiveBayesClassifier<TFeature, TClass>> implementation

        public bool Equals(NaiveBayesClassifier<TFeature, TClass> other)
        {
            if (other == null)
                throw new ArgumentNullException("other");

            return Prior.DictionaryEqual(other.Prior);
        }

        #endregion

        #region Object overriden

        public override int GetHashCode()
        {
            return Prior.OrderBy(x => x.Key).Aggregate(1, (hashCode, x) => hashCode ^ x.Key.GetHashCode());
        }

        public override bool Equals(object other)
        {
            if (other == null)
                throw new ArgumentNullException("other");

            if (!(other is NaiveBayesClassifier<TFeature, TClass>))
                throw new InvalidCastException("The 'other' argument is not a NaiveBayesClassifier<TFeature, TClass> object.");

            return Equals(other as NaiveBayesClassifier<TFeature, TClass>);
        }

        public override string ToString()
        {
            return String.Format("Prior: {0}, Likelihood: {1}", Prior,
                                 String.Join(", ", Likelihood.Select(x => String.Format("{{{0}, {1}}}", x.Key, x.Value))));
        }

        #endregion
    }
}
